# coding=utf-8
"""
author: jamon
"""
from __future__ import annotations

import asyncio
import time
import weakref
from asyncio import Future
from enum import IntEnum
from typing import Optional, Any, TYPE_CHECKING

import ujson
import websockets

from client.net.msg_protocol import MsgProtocol
from obespoir.share.espoir_log import global_log as logger
from package.universe import EsException

if TYPE_CHECKING:
    from client.net.conn_manager import NetConnectManager


class ConnectStatus(IntEnum):
    NOT_CONNECT = 0  # 未连接
    CONNECTED = 1  # 已连接
    CLOSED = 2  # 已断开


def get_url(host, port):
    return f"ws://{host}:{port}"


class MsgIndex(object):
    last_time = int(time.time())
    last_index = 1

    @classmethod
    def next(cls) -> str:
        now = int(time.time())

        if cls.last_time == now:
            cls.last_index += 1
        else:
            cls.last_time = now
            cls.last_index = 1

        return "{}#{}".format(cls.last_time, cls.last_index)


class WebSocketService(object):

    def __init__(self, host: str, port: int, conn_mgr: NetConnectManager):
        self.ws_host: str = host
        self.ws_port: int = port

        self.ws_url = get_url(host, port)
        self.conn_mgr: NetConnectManager = weakref.proxy(conn_mgr)

        self.cur_status: ConnectStatus = ConnectStatus.NOT_CONNECT  # 连接状态
        self.conn: Optional[websockets.WebSocketClientProtocol] = None
        self.protocol: MsgProtocol = MsgProtocol()

        # 剩余自动重连的次数
        self.reconnect_times: int = 5

        # 回调future
        self.future_dict: dict[str, Future] = {}

        self.loop = asyncio.get_running_loop()

        self.conn_callback = None

        self.handle_message = None

    async def run(self):
        logger.info("start connect server", ws_url=self.ws_url)
        async with websockets.connect(self.ws_url) as websocket:
            logger.info("connected success", ws_url=self.ws_url)
            self.cur_status = ConnectStatus.CONNECTED
            self.conn = websocket
            if self.conn_callback:
                self.conn_callback()

            while ConnectStatus.CONNECTED == self.cur_status:
                try:
                    recv_data = await self.conn.recv()
                except websockets.ConnectionClosed:
                    self.cur_status = ConnectStatus.CLOSED
                else:
                    await self.handle_receive(recv_data)
        if ConnectStatus.CLOSED == self.cur_status:
            self.reconnect_times -= 1
            if 0 >= self.reconnect_times:
                logger.warn("reconnect over max times", ws_url=self.ws_url)
                self.conn_mgr.reset_ws_conn()
                return
            await self.run()

    async def handle_receive(self, recv_data: bytes) -> Any:
        logger.info("start", recv_data=recv_data)
        unpack_msg = self.protocol.unpack(recv_data)
        unpack_data = unpack_msg["data"]
        msg_id = unpack_msg["command"]
        if isinstance(unpack_data, (str, bytes)):
            unpack_data = ujson.loads(unpack_data)
        seq = unpack_data.get("seq", None)
        session_id = unpack_data.get("session_id", None)
        data = unpack_data.get("data", unpack_data)
        if seq and seq in self.future_dict:
            self.future_dict[seq].set_result(data)
            del self.future_dict[seq]
        logger.info("end", msg_id=msg_id, seq=seq, unpack_data=unpack_data,
                    session_id=session_id)
        asyncio.create_task(self.handle_message(msg_id, seq, session_id,
                                                unpack_data))

    async def handle_send(self, command_id: int, data: dict) -> Any:
        begin_time = time.time()
        while not self.conn:
            await asyncio.sleep(0.1)
            cur_time = time.time()
            if 3 < cur_time - begin_time:
                if not self.conn:
                    raise EsException("ConnectionObjectError",
                                      ws_url=self.ws_url)

        seq = MsgIndex.next()
        to_pack_data = {"data": data, "seq": seq}
        packed_data = self.protocol.pack(ujson.dumps(to_pack_data), command_id)
        logger.info("start", command_id=command_id, data=data, seq=seq)
        wait_future: Future = self.loop.create_future()
        await self.conn.send(packed_data)
        self.future_dict[seq] = wait_future
        logger.info("end", command_id=command_id, seq=seq)
        return await wait_future
